import re
from math import sqrt
from numpy import linspace, dot, complex128, ones, reshape,\
                  array, eye, kron, append, zeros, std, mean
from numpy.random import rand, randn
from c_traj import get_trajectory
#from get_trajectory_test import get_trajectory
import datetime as dt
import os
from os.path import isfile
import cPickle as pickle

def current_timestamp():
    present = dt.datetime.now()
    ymd = '_'.join(map(str, [present.year, present.month, present.day]))
    wall_time = '_'+str(present.hour) + ':' + str(present.minute).rjust(2,'0')+ ':' + str(present.second).rjust(2,'0')
    return ymd + wall_time

def prep_dict_indep():
    """
    Place to dump messy code which sets the problem-independent values 
    of the input_dict. This ensures that unnecessary kronecker products
    are only evaluated once, not every time the objective
    function is called.
    """

    #Time-independent part of Liouvillian
    sigmaZ = array([[1, 0], [0, - 1]])
    sigmaM = array([[0, 1], [0, 0]])
    smalleye = eye(2)
    fulleye = eye(8)

    # The dephasing operator for qubit 1
    Ad1 = reduce(kron,(sigmaZ, smalleye, smalleye))
    Ld1 = kron(Ad1, Ad1.T) - dot(0.5, kron(dot(
        Ad1.T, Ad1), fulleye)) - dot(0.5, kron(fulleye, dot(Ad1.T, Ad1).T))

    # The dephasing operator for qubit 2
    Ad2 = reduce(kron,(smalleye, sigmaZ, smalleye))
    Ld2 = kron(Ad2, Ad2.T) - dot(0.5, kron(dot(
        Ad2.T, Ad2), fulleye)) - dot(0.5, kron(fulleye, dot(Ad2.T, Ad2).T))

    # The dephasing operator for qubit 3
    Ad3 = reduce(kron,(smalleye, smalleye, sigmaZ))
    Ld3 = kron(Ad3, Ad3.T) - dot(0.5, kron(dot(
        Ad3.T, Ad3), fulleye)) - dot(0.5, kron(fulleye, dot(Ad3.T, Ad3).T))

    # The relaxation operator for qubit 1
    Ar1 = reduce(kron,(sigmaM, smalleye, smalleye))
    Ar1dagger = Ar1.T
    Lr1 = kron(Ar1, Ar1dagger.T) - dot(0.5, kron(dot(
        Ar1dagger, Ar1), fulleye)) - dot(0.5, kron(fulleye, dot(Ar1dagger, Ar1).T))

    # The relaxation operator for qubit 2
    Ar2 = reduce(kron,(smalleye, sigmaM, smalleye))
    Ar2dagger = Ar2.T
    Lr2 = kron(Ar2, Ar2dagger.T) - dot(0.5, kron(dot(
        Ar2dagger, Ar2), fulleye)) - dot(0.5, kron(fulleye, dot(Ar2dagger, Ar2).T))

    # The relaxation operator for qubit 3
    Ar3 = reduce(kron,(smalleye, smalleye, sigmaM))
    Ar3dagger = Ar3.T
    Lr3 = kron(Ar3, Ar3dagger.T) - dot(0.5, kron(dot(
        Ar3dagger, Ar3), fulleye)) - dot(0.5, kron(fulleye, dot(Ar3dagger, Ar3).T))

    #Ideal Plus/Minus density matrices
    ideal_p = 0.5 * array([1,0,0,1,0,1,1,0], dtype=complex128).T
    ideal_p = kron(ideal_p, ideal_p)
    ideal_m = 0.5 * array([0,1,1,0,1,0,0,1], dtype=complex128).T
    ideal_m = kron(ideal_m, ideal_m)
  
    initstate = 1./sqrt(8.) * ones((8, 1), dtype='complex128')
    initMat = initstate * initstate.T
    init = reshape(initMat.T, 64, 1)

    indep_dict = {'dim': 8, 'chi': 1., 'kappa1': 2., 'kappa2': 2., 
                    'Delta_a': sqrt(3.), 'Delta_b': -sqrt(3.),
                    'dt': 0.005, 'init': init, 'ideal_p': ideal_p,
                    'ideal_m': ideal_m, 'Ld1': Ld1, 'Ld2': Ld2,
                    'Ld3': Ld3, 'Lr1': Lr1, 'Lr2': Lr2, 'Lr3': Lr3}

    return indep_dict

def prep_dict_dep(dict_so_far, params):
    """
    Prepares a dictionary for input to get_trajectory, based on the
    parameters of a pre-existing dictionary.

    This function takes the time to be optimized to be the drive
    time, so that the measurement turns off at the same time as the
    drive, and the simulation continues for another 2/chi to capture
    the ring-down. 
    """
    
    if (probject._gamma_d == 0) and (probject._gamma_r == 0):
        t_drive = probject._fixed_time
        pass_params = params
    else:
        t_drive = params[-1]
        pass_params = params[:-1]
    
    tmeas_off = t_drive

    #Decay rates, given by master decay rates gamma_d and gamma_r

    #input_dict['T'] = distortion_fudge_factor * input_dict['tmeas_off'] # final time
    T = tmeas_off # only drive time here

    #Cavity input, given by pulse_func:
    dt = dict_so_far['dt']
    N = int(round(T / dt))
    tvec = linspace(0., T, N)
    epsilon = probject._pulse_func(pass_params, tvec)

    # The time independent deterministic part of Liouvillian, affected by kappa
    #input_dict['D1'] = (dot(dot(0.5, input_dict['gamma_d1']), Ld1) + \
    #                    dot(dot(0.5, input_dict['gamma_d2']), Ld2) + \
    #                    dot(dot(0.5, input_dict['gamma_d3']), Ld3) + \
    #                    dot(dot(0, (dot((g / Delta_q) ** 2, (input_dict['kappa1'] + input_dict['kappa2'])) + \
    #                        input_dict['gamma_1'])), (Lr1 + Lr2 + Lr3))).astype(complex128)
    
    gamma_d = probject._gamma_d
    gamma_r = probject._gamma_r
    Ld1 = probject._input_dict['Ld1']; Ld2 = probject._input_dict['Ld2']
    Ld3 = probject._input_dict['Ld3']
    Lr1 = probject._input_dict['Lr1']; Lr2 = probject._input_dict['Lr2']
    Lr3 = probject._input_dict['Lr3']
    kappa1 = probject._input_dict['kappa1']; kappa2 = probject._input_dict['kappa2']
    D1 = (0.5 * gamma_d * Ld1 + 0.5 * gamma_d * Ld2 + 0.5 * gamma_d * Ld3 + 
            0 * (1./1600. * (kappa1 + kappa2) + gamma_r ) * (Lr1 + Lr2 + Lr3)).astype(complex128)

    dep_dict = {'D1': D1, 'eta': probject._eta, 'tmeas_on': 0.,
                'tmeas_off': tmeas_off, 'epsilon': epsilon, 'N': N, 'T': T}

    return dep_dict

def gaussian_pulse(h_offset, h_peak, st_dev, times):
    """
    Raised windowed Gaussian for locally-optimal driving
    """
    from numpy import amax, exp
    mean = amax(times)/2.
    gaus_array = h_peak * exp(-(times - mean)**2/(2. * st_dev ** 2)) + h_offset
    return gaus_array

def is_timestamped(direc_str):
    """
    Returns True if a MatchObject can be created from a specific
    pattern matching a timestamp, False if None is created.    
    """
    #matches four digits and 1 or 2 digits separated appropriately.
    #May 12 2014 two timestamp formats now used due to faster code. 
    stmp_ptrn = '[0-9]{4}_[0-9]{1,2}_[0-9]{1,2}_[0-9]{1,2}:[0-9]{1,2}$'
    stmp_s_ptrn = '[0-9]{4}_[0-9]{1,2}_[0-9]{1,2}_[0-9]{1,2}:[0-9]{1,2}:[0-9]{1,2}$'
    return (bool(re.match(stmp_ptrn, direc_str)) or bool(re.match(stmp_s_ptrn, direc_str)))

def confidence_interval(data, p):
    sort_data = sorted(list(data))
    len_outside = int(0.5 * (1. - p) * len(sort_data))
    #print p, len(sort_data)
    return max(sort_data[:len_outside]), min(sort_data[-len_outside:])

def is_correct(s_value, overlap_p, overlap_m):
    if s_value > 0:
        return True if overlap_p > overlap_m else False
    elif s_value < 0:
        return True if overlap_p < overlap_m else False
    elif s_value == 0:
        return True if overlap_p == overlap_m else False

def meas_fid(s_lst, p_lst, m_lst, alpha=0.68, detection='real'):
    if detection == 'real':
        fid_list = [
                    sqrt(max(overlap_p, overlap_m))
                    if is_correct(s_value, overlap_p, overlap_m) else
                    sqrt(min(overlap_p, overlap_m)) 
                    for s_value, overlap_p, overlap_m in zip(s_lst, p_lst, m_lst)
                    ]
    elif detection == 'ideal':
        fid_list = [
                sqrt(max(overlap_p, overlap_m))
                for overlap_p, overlap_m in zip(p_lst, m_lst)
                ]
    return sum(fid_list) / len(fid_list), confidence_interval(fid_list, alpha)

def pulse_fitness(input_dict, detection='real'):
    output_dict = get_trajectory(input_dict)
    if detection=='real':
        if output_dict['s_filtered'].real > 0:
            return (sqrt(output_dict['overlap_p']),)
        elif output_dict['s_filtered'].real < 0.:
            return (sqrt(output_dict['overlap_m']),)
        else:
            #might be incorrect, but s is never 0.
            return (0.5 * (sqrt(output_dict['overlap_m']) + 
                            sqrt(output_dict['overlap_p'])),)
    elif detection=='ideal':
        return (sqrt(max([output_dict['overlap_m'],output_dict['overlap_p']])),)

traj_fid = lambda input_dict, detection: pulse_fitness(input_dict, detection)[0]

def stoc_eval(fun, args, tol=5*10**-3, alpha=0.68, maxiter=100, iter_check=False):
    """
    Evaluates a stochastic function (one which returns an answer 
    partially-dependent on random parameters) by looping until the
    relative standard deviation of the mean is less than some
    tolerance. Returns the mean and a confidence interval which contains 
    68% of the produced values by default. Has an iteration cap of
    10000.
    """
    #Running calculation of mean and standard deviation from
    #wikipedia.org/wiki/Standard_deviation#Rapid_calculation_methods
    mean_val = 0
    old_mean_val = 0
    stdev = 0
    k = 0
    #ensure first ten iterations to avoid early exit
    rel_std_mean = 1
    data=[]
    while ((rel_std_mean > tol) or (k < 10)) and (k < maxiter):
        k += 1; sample = fun(*args)
        data.append(sample)
        old_mean_val = mean_val
        mean_val += (sample - mean_val) / k
        stdev += (sample - old_mean_val) * (sample - mean_val)
        rel_std_mean = sqrt(stdev) / (k * mean_val)
        #print k, sample
    if iter_check:
        return mean_val, confidence_interval(data, alpha), k
    else:
        return mean_val, confidence_interval(data, alpha)

def test_fun(x):
    """
    Simple test function for stochastic optimization, returns
    ((x - 2) ** 2 + 10) * (1. + 0.03 * randn())
    """
    actual_value = (x - 2.) ** 2. + 10.
    return actual_value * (1. + 0.03 * randn())

onlyfiles = lambda dir_name: [ f for f in os.listdir(dir_name) if isfile('/'.join([dir_name, f]))]

def dir_iter_check(dir_name, maxiter = 90000):
    for phil in onlyfiles(dir_name):
        if phil != "README":
            with open(dir_name + '/' + phil,'r') as phile:
                dicto = pickle.load(phile)
                if dicto['iters'] > maxiter:
                    print dicto['iters']
























